from flat_quant.utils.model_utils import (ModelStructureBridge, 
                                            AttnNormLinearPair, 
                                            AttnLinearLinearPair, 
                                            MLPNormLinearPair, 
                                            MLPLinearLinearPair)


class LlamaStructureBridge(ModelStructureBridge):
    """Llama模型结构桥接器"""
    def analyze_structure(self):
        # 分析Llama模型结构，添加所有相关结构对
        attn_norm_linear_names = ["self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj"]
        attn_linear_linear_names = ["self_attn.o_proj"]
        mlp_norm_linear_names = ["mlp.gate_proj", "mlp.up_proj"]
        mlp_linear_linear_names = ["mlp.down_proj"]
        for name, module in self.model.named_modules():
            if "input_layernorm" in name:
                linears = []
                for linear_name in attn_norm_linear_names:
                    linears.append(name.replace("input_layernorm", linear_name))
                prefix_name = '.'.join(name.split('.')[:3])
                self.register_structure_pair(AttnNormLinearPair(self.config, name, linears, prefix_name))
            if "self_attn.v_proj" in name:
                linears = []
                for linear_name in attn_linear_linear_names:
                    linears.append(name.replace("self_attn.v_proj", linear_name))
                prefix_name = '.'.join(name.split('.')[:3])
                self.register_structure_pair(AttnLinearLinearPair(self.config, name, linears, prefix_name))
            if "post_attention_layernorm" in name:
                linears = []
                for linear_name in mlp_norm_linear_names:
                    linears.append(name.replace("post_attention_layernorm", linear_name))
                prefix_name = '.'.join(name.split('.')[:3])
                self.register_structure_pair(MLPNormLinearPair(self.config, name, linears, prefix_name))
            if "mlp.up_proj" in name:
                linears = []
                for linear_name in mlp_linear_linear_names:
                    linears.append(name.replace("mlp.up_proj", linear_name))
                prefix_name = '.'.join(name.split('.')[:3])
                self.register_structure_pair(MLPLinearLinearPair(self.config, name, linears, prefix_name))
        self._layers_name = "model.layers"